# Some basic setup:
# Setup detectron2 logger
import detectron2
from detectron2.utils.logger import setup_logger
setup_logger()

# import some common libraries
import numpy as np
import os, json, cv2, random
import matplotlib.pyplot as plt  # 替换google.colab的cv2_imshow

# import some common detectron2 utilities
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog, DatasetCatalog

# 检查CUDA是否可用
import torch
print("CUDA available:", torch.cuda.is_available())


# 读取图像
im = cv2.imread("./input/person9_4.45.jpg")
if im is None:
    raise FileNotFoundError("无法加载图像，请检查input.jpg是否存在")

# 显示图像 - 替换cv2_imshow
def cv2_imshow(image):
    """替代Google Colab的cv2_imshow函数"""
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    plt.imshow(image)
    plt.axis('off')
    plt.show()

cv2_imshow(im)

# 配置模型
cfg = get_cfg()
# 添加项目特定配置
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5  # 设置阈值

# 设置GPU设备
cfg.MODEL.DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# 从model zoo加载预训练模型权重
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")

# 创建预测器
predictor = DefaultPredictor(cfg)

# 进行预测
outputs = predictor(im)

# 可视化结果
v = Visualizer(im[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.0)     #此处的scale代表缩放因子
out = v.draw_instance_predictions(outputs["instances"].to("cpu"))

# 转换为numpy数组并确保可写
vis_image = np.array(out.get_image())[:, :, ::-1].copy()

# 添加编号
instances = outputs["instances"].to("cpu")
boxes = instances.pred_boxes.tensor.cpu().numpy()  # 获取所有框的坐标
for i, box in enumerate(boxes):
    cv2.putText(vis_image, str(i), (int(box[0]), int(box[1])-10),
                cv2.FONT_HERSHEY_SIMPLEX, 3, (0, 0, 255), 4)

cv2_imshow(vis_image)

# 让用户输入要计算距离的物体编号
while True:
    try:
        selected_id = int(input("\n请输入要计算距离的物体编号: "))
        if 0 <= selected_id < len(boxes):
            break
        else:
            print(f"错误: 编号必须在0-{len(boxes)-1}之间")
    except ValueError:
        print("错误: 请输入有效的数字编号")

# 距离检测
'''
#print(outputs["instances"].pred_classes)
#print(outputs["instances"].pred_boxes)
#print(outputs["instances"].scores)
#print(MetadataCatalog.get(cfg.DATASETS.TRAIN[0]))

for data in outputs["instances"].pred_classes:
  num = data.item();
  print(MetadataCatalog.get(cfg.DATASETS.TRAIN[0]).thing_classes[num]);
'''
delta_y = outputs["instances"].pred_boxes.tensor;
delta_y1 = delta_y[selected_id][3] - delta_y[selected_id][1]
delta_y1 = delta_y1.item()        #识别框高度的像素值
print(delta_y)
print(delta_y1)

# 获取高度和宽度
height_px, width_px = im.shape[:2]
print(f"宽度(像素): {width_px}, 高度(像素): {height_px}")
mi_dis_mm = 5.59      #相机焦距，单位mm
tensor_w = 3.59       #摄像头传感器模块的宽度，单位mm
tensor_h = 7.4        #摄像头传感器模块的高度，单位mm
height_phy = 1.7     #物体的实际高度，单位mm
width_phy = 0.3       #物体的实际宽度，单位mm


if height_px > width_px:    #竖屏拍摄
    print("竖屏拍摄");
    mi_dis_px = mi_dis_mm * height_px / tensor_h;      #相机焦距，单位：像素
    distance = height_phy * mi_dis_px / delta_y1;      #计算所得距离
else:
    print("横屏拍摄");
    mi_dis_px = mi_dis_mm * height_px / tensor_w;    #横屏拍摄
    distance = height_phy * mi_dis_px / delta_y1;

print("distance = ", distance)

# 将距离写入 result.txt 文件
with open("result.txt", "w") as f:
    f.write(f"{distance}\n")  # 你可以根据需要调整格式

# 推送到 Gitee
import git
import os


# 检查是否是 Git 仓库，如果不是则初始化
if not os.path.exists(os.path.join(os.getcwd(), ".git")):
    repo = git.Repo.init(os.getcwd())
    print("✅ 初始化 Git 仓库成功")
else:
    repo = git.Repo(os.getcwd())

repo = git.Repo(os.getcwd())
# 设置远程仓库（Gitee）
origin_url = "git@gitee.com:whynowhy/my_detectron2.git"  # 替换成你的仓库地址
if "origin" not in repo.remotes:
    repo.create_remote("origin", origin_url)
else:
    repo.remote("origin").set_url(origin_url)

print("✅ Git 仓库已配置，远程地址:", origin_url)

# 确保当前目录是 Git 仓库
repo = git.Repo(os.getcwd())

try:
    # 添加文件更改
    repo.git.add("result.txt")
    repo.git.add("vision.py")
    #repo.git.add("input")

    # 提交更改
    repo.git.commit("-m", "Update distance result")

    # 拉取远程更改（避免冲突）
    repo.git.pull("origin", "master")  # 拉取远程 master 分支
    # 推送到 Gitee（SSH 方式）
    origin = repo.remote(name="origin")
    origin.push(force=True)     #强制推送，如果不加 "force=True" ,当gitee仓库有更改，如Readme文件更改时，本地文件不能被推送到仓库

    print("✅ Successfully pushed to Gitee via SSH!")
except Exception as e:
    print(f"❌ Error pushing to Gitee: {e}")